//
// Created by wangwanlong on 2024/11/5.
//

#ifndef MKAHYPAR_RANDOM_INITIAL_PARTITIONER_H
#define MKAHYPAR_RANDOM_INITIAL_PARTITIONER_H

#include <queue>
#include "mkahypar/partition/initial_partitioning/i_initial_partitioner.h"
#include "mkahypar/partition/initial_partitioning/initial_partitioner_base.h"
#include "mkahypar/partition/refinement/policies/hop_penalty_policy.h"
#include "mkahypar/partition/refinement/policies/block_penalty_policy.h"
#include "mkahypar/partition/refinement/policies/edge_penalty_policy.h"

//基于随机的搜索放置
namespace mkahypar {
  class Search0InitialPartitioner : public IInitialPartitioner,
                                    public InitialPartitionerBase<Search0InitialPartitioner> {
    using Base = InitialPartitionerBase<Search0InitialPartitioner>;
  public:
    Search0InitialPartitioner(Hypergraph &hypergraph, TargetGraph &targetGraph) :
            Base(hypergraph, targetGraph), _is_critical_node(hypergraph.initialNumNodes()) {
    }

    ~Search0InitialPartitioner() override = default;


    void initialPartition() {
      //识别关键顶点:主要思想考虑一个fpga资源占比和总互联量两个指标
      std::vector<std::pair<double, HypernodeID>> critical_nodes;
      critical_nodes.reserve(_hg._num_hypernodes);

      for (auto hn: _hg.nodes()) {
        double p = getMaxResPercent(hn);
        GainType he_weight = 0;
        for (auto he: _hg.incidentEdges(hn))
          he_weight += _hg.edgeWeight(he);
        critical_nodes.emplace_back(p * he_weight, hn);
      }
      std::sort(critical_nodes.begin(), critical_nodes.end());

      int idx = critical_nodes.size() * Context::getInstance().critical_node_factor;
      idx = std::max(1, idx);//保证至少一个关键顶点
      std::queue<HypernodeID> que;
      ds::FastResetFlagArray<>in_que(_hg.initialNumNodes());

      for (int i = 0; i < idx; ++i) {
        _is_critical_node.set(critical_nodes[i].second, true);//标记关键顶点
        que.emplace(critical_nodes[i].second);
        in_que.set(critical_nodes[i].second,true);
      }
//      LOG<<"start place large nodes:";


      while (!que.empty()) {
        auto hn = que.front();
        que.pop();
        //1.给hn找一个最好的位置
        Metrics best_gain;
        PartitionID best_part = Hypergraph::kInvalidPartition;
        best_gain.obj = std::numeric_limits<GainType>::min();
        //搜索最佳移动分区
        findBestMove(hn, best_gain, best_part);
        //尝试移动道best_part
        placeNode(hn, best_part);
//        LOG<<"assign"<<hn<<"to"<<best_part;

        //2.扩展hn的邻居大节点
        for(auto he : _hg.incidentEdges(hn)){
          for(auto pin : _hg.pins(he)){
            if(_hg.partID(pin) == Hypergraph::kInvalidPartition && _is_critical_node[pin] &&
            !in_que[pin]){
              //没有被放置的关键顶点
              que.emplace(pin);
              in_que.set(pin, true);
            }
          }
        }
      }
      //3.放置剩余的小节点
      std::vector<HypernodeID>small_nodes;
      for(auto hn : _hg.nodes()){
        if(_hg.partID(hn) == Hypergraph::kInvalidPartition)
          small_nodes.emplace_back(hn);
      }
//      LOG<<"start place small nodes:";
      Randomize::getInstance().shuffleVector(small_nodes,small_nodes.size());
      for(auto hn : small_nodes){
        Metrics best_gain;
        PartitionID best_part = Hypergraph::kInvalidPartition;
        best_gain.obj = std::numeric_limits<GainType>::min();
        //搜索最佳移动分区
        findBestMove(hn, best_gain, best_part);
//        DBG<<"assign"<<hn<<"to"<<best_part;
        //尝试移动道best_part
        placeNode(hn, best_part);
      }

      //4.初始化
      for (HypernodeID hn = 0; hn < _hg.initialNumNodes(); ++hn)
        _hg.hypernode(hn).num_incident_cut_hes = 0;
      for (const HyperedgeID &he: _hg.edges()) {
        if (_hg.connectivity(he) > 1) {
          for (const HypernodeID &pin: _hg.pins(he)) {
            ++_hg.hypernode(pin).num_incident_cut_hes;
          }
        }
      }

    }

    //计算节点最大的资源的占比
    double getMaxResPercent(HypernodeID v) {
      auto &fpga_res = _topo._fpga[0]._bound_res;
      double p = 0;
      for (int i = 0; i < _hg.hypernode(v).weight().size(); ++i) {
        if (_hg.hypernode(v).weight()[i] != 0) {
          double x = static_cast<double>(_hg.hypernode(v).weight()[i]) /
                     static_cast<double>(fpga_res[i]);
          p = std::max(p, x);
        }
      }
      return p;
    }


    void findBestMove(HypernodeID v, Metrics &best_gain, PartitionID &best_part) {
      PartitionID from = _hg.partID(v);
      for (PartitionID id = 0; id < _hg._k; ++id) {
        if (id == from)
          continue;
        Metrics delta_gain;
        evaluateNodeMoveGain(v, from, id, delta_gain);
        if (delta_gain.obj > best_gain.obj) {
          best_part = id;
          best_gain = delta_gain;
        }
      }
    }


    //评估节点移动收益,其中from是非法分区,to是合法分区
    template<class BlockPenaltyPolicy = NormalBlockPenalty,
            class EdgePenaltyPolicy = NormalEdgePenaltyPolicy,
            class HopPenaltyPolicy = NormalHopPenaltyPolicy>
    void evaluateNodeMoveGain(HypernodeID v, PartitionID from, PartitionID to, Metrics &delta_gain) {
      if (from != Hypergraph::kInvalidPartition) {
        LOG<<v<<from;
        ERROREXIT("error partition!");
      }

      //计算资源的变化量
      delta_gain.block_penalty += BlockPenaltyPolicy::computeAdd(_hg.nodeWeight(v),
                                                                 _topo._fpga[to]._used_res,
                                                                 _topo._fpga[to]._bound_res);
      std::vector<GainType> line_delta(_hg._k);
      //计算路由和超载
      for (auto &he: _hg.incidentEdges(v)) {
        auto &edge = _hg.hyperedge(he);
        auto he_weight = _hg.edgeWeight(he);
        auto source = _hg.hyperedge(he).source();
        HypernodeID source_id = edge.source();
        bool only_one_pin_in_to_part = _hg.pinCountInPart(he, to) == 0;
        if (only_one_pin_in_to_part) {
          if (_hg.connectivity(he) == 1) {
            //内部边 -> 边界边，to和唯一连接的part新增互联量
            line_delta[to] += he_weight;
            line_delta[_hg.hyperedge(he).connectivity.getOnlyPos()] += he_weight;
          } else if (_hg.connectivity(he) > 1) {
            //边界边 -> 边界边，新增to的互联量
            line_delta[to] += he_weight;
          }
        }
        //计算路由变化
        if (v == source) {
          for (PartitionID pid = 0; pid < _hg._k; ++pid) {
            GainType count = _hg.routeCountInPart(he, pid);
            if (count > 0) {
              delta_gain.hop_penalty -= HopPenaltyPolicy::compute(he_weight, _topo._distance[to][pid]) * count;
            }
          }
        } else {
          PartitionID source_part = _hg.partID(source);
          if (source_part == Hypergraph::kInvalidPartition)
            continue;
          GainType count = _hg.tarRouteCount(he, v);
          delta_gain.hop_penalty -= HopPenaltyPolicy::compute(he_weight, _topo._distance[source_part][to]) * count;
        }
      }
      //计算互联的变化量
      for (PartitionID id = 0; id < _hg._k; ++id) {
        if (line_delta[id] > 0)
          delta_gain.edge_penalty += EdgePenaltyPolicy::computeDelta(_topo._fpga[id]._used_line,
                                                                     _topo._fpga[id]._used_line + line_delta[id],
                                                                     _topo._fpga[id]._bound_line);
      }
      delta_gain.updateObj();

    }

    //放置节点hn到to,hn必须是invalid part
    void placeNode(HypernodeID hn, PartitionID to){
      if(_hg.partID(hn) != Hypergraph::kInvalidPartition)
        ERROREXIT("error part!");
      //资源量更正
      _topo._fpga[to]._used_res += _hg.hypernode(hn).weight();
      //更改节点分区
      _hg.hypernode(hn).part_id = to;
      for (const HyperedgeID &he: _hg.incidentEdges(hn)){
        HypernodeID source = _hg.hyperedge(he).source();
        HyperedgeWeight he_weight = _hg.edgeWeight(he);
        //更新pin_in_part
        PartitionID only_part = _hg.hyperedge(he).connectivity.getOnlyPos();
        bool only_one_pin_in_to_part = _hg.incrementPinCountInPart(he, to);
        //更新target_part_route_count
        if(hn != source) {
          _hg.incrementRouteCountInPart(he, to, hn);
        }
        //更新互联量
        if(only_one_pin_in_to_part){
          if (_hg.connectivity(he) == 2) {
            //内部边 -> 边界边，to和唯一连接的part新增互联量
            _topo._fpga[to]._used_line += he_weight;
            _topo._fpga[only_part]._used_line += he_weight;
          } else if (_hg.connectivity(he) > 2) {
            //边界边 -> 边界边，新增to的互联量
            _topo._fpga[to]._used_line += he_weight;
          }
        }
      }

    }

    void partitionImpl() override final {
      Base::singleRunsInitialPartitioning();
    }

    ds::FastResetFlagArray<> _is_critical_node;
  };
}
#endif //MKAHYPAR_RANDOM_INITIAL_PARTITIONER_H
